package com.ls.server.mcp.client;

import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.client.transport.WebFluxSseClientTransport;
import org.springframework.ai.mcp.client.autoconfigure.NamedClientMcpTransport;
import org.springframework.ai.mcp.client.autoconfigure.properties.McpSseClientProperties;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;

import java.net.http.HttpClient;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;


@Component
public class SseWebFluxTransport {

    ObjectMapper objectMapper = new ObjectMapper();

    /**
     *
     * 功能描述: Creates a list of WebFlux-based SSE transports for MCP communication.
     *
     * @param connections
     * @return java.util.List<org.springframework.ai.mcp.client.autoconfigure.NamedClientMcpTransport>
     * @author 刘帅
     * @date 2025-05-08 11:29
     * @Version:
     */
    public List<NamedClientMcpTransport> webFluxClientTransports(Map<String, McpSseClientProperties.SseParameters> connections) {
        var webClientBuilderTemplate = WebClient.builder();
        var sseTransports = new ArrayList();
        for (Map.Entry<String, McpSseClientProperties.SseParameters> serverParameters : connections.entrySet()) {
            var webClientBuilder = webClientBuilderTemplate.clone().baseUrl(serverParameters.getValue().url());
            var transport = new WebFluxSseClientTransport(webClientBuilder, objectMapper);
            sseTransports.add(new NamedClientMcpTransport(serverParameters.getKey(), transport));
        }
        return sseTransports;
    }

    public List<NamedClientMcpTransport> mcpHttpClientTransports(Map<String, McpSseClientProperties.SseParameters> connections) {
        List<NamedClientMcpTransport> sseTransports = new ArrayList();
        connections.forEach((key, value) -> {
            HttpClientSseClientTransport transport = new HttpClientSseClientTransport(HttpClient.newBuilder(), value.url(), objectMapper);
            sseTransports.add(new NamedClientMcpTransport(key, transport));
        });
        return sseTransports;
    }
}
